package com.example.memoizer;

import java.math.BigInteger;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;

public class Memoizer<A, V> implements Computable<A, V> {
	private static final int EXPIRATION_TIME = 3; // 3 seconds
	
	private final Computable<A, V> computer;
	private final ConcurrentMap<A, FutureTask<V>> cache = new ConcurrentHashMap<>();
	private final ConcurrentMap<A, Long> execTime = new ConcurrentHashMap<>();

	public Memoizer(Computable<A, V> computer) {
		this.computer = computer;
	}

	@Override
	public V compute(final A arg) throws InterruptedException {
		while (true) {
			Long time = execTime.get(arg);
			boolean expired;
			if (time == null)
				expired = false;
			else
				expired = timestamp() - execTime.get(arg) > EXPIRATION_TIME;
			
			if (expired)
				cache.remove(arg);
			
			Future<V> f = cache.get(arg);
			
			if (f == null) {
				FutureTask<V> ft = new FutureTask<>(new Callable<V>() {
					@Override
					public V call() throws Exception {
						V result = computer.compute(arg);
						execTime.put(arg, timestamp());
						return result;
					}
				});

				f = cache.putIfAbsent(arg, ft);
				if (f == null) {
					f = ft;
					ft.run();
				}
			}

			try {
				return f.get();
			} catch (ExecutionException e) {
				cache.remove(arg);
				throw new RuntimeException("Computation failed", e);
			} catch (CancellationException e) {
				cache.remove(arg);
			}
		}
	}
	
	private long timestamp() {
		return System.currentTimeMillis() / 1000;
	}
}

class ExpensiveComputation implements Computable<String, BigInteger> {

	@Override
	public BigInteger compute(String arg) throws InterruptedException {
		System.out.println("computing");
		return new BigInteger(arg);
	}
}